{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Dash Summarize Colab Demo.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyMNRCtZ4l+aaOF2O9TB+rzK"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8n6CYTIja4hq",
        "colab_type": "text"
      },
      "source": [
        "To start this Jupyter Dash app, please run all the cells below. Then, click on the **temporary** URL at the end of the last cell to open the app."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XSlKBZjYotV1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install -q jupyter-dash==0.3.0rc1 dash-bootstrap-components transformers"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XTZjvYnkn45v",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import time\n",
        "\n",
        "import dash\n",
        "import dash_html_components as html\n",
        "import dash_core_components as dcc\n",
        "import dash_bootstrap_components as dbc\n",
        "from dash.dependencies import Input, Output, State\n",
        "from jupyter_dash import JupyterDash\n",
        "from transformers import BartTokenizer, BartForConditionalGeneration\n",
        "import torch"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aN5EHFXSoDTv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "print(f\"Device: {device}\")\n",
        "\n",
        "# Load Model\n",
        "pretrained = \"sshleifer/distilbart-xsum-12-6\"\n",
        "model = BartForConditionalGeneration.from_pretrained(pretrained)\n",
        "tokenizer = BartTokenizer.from_pretrained(pretrained)\n",
        "\n",
        "# Switch to cuda, eval mode, and FP16 for faster inference\n",
        "if device == \"cuda\":\n",
        "    model = model.half()\n",
        "model.to(device)\n",
        "model.eval();"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z2EdDQbqpgy1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Define app\n",
        "app = JupyterDash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])\n",
        "server = app.server\n",
        "\n",
        "controls = dbc.Card(\n",
        "    [\n",
        "        dbc.FormGroup(\n",
        "            [\n",
        "                dbc.Label(\"Output Length (# Tokens)\"),\n",
        "                dcc.Slider(\n",
        "                    id=\"max-length\",\n",
        "                    min=10,\n",
        "                    max=50,\n",
        "                    value=30,\n",
        "                    marks={i: str(i) for i in range(10, 51, 10)},\n",
        "                ),\n",
        "            ]\n",
        "        ),\n",
        "        dbc.FormGroup(\n",
        "            [\n",
        "                dbc.Label(\"Beam Size\"),\n",
        "                dcc.Slider(\n",
        "                    id=\"num-beams\",\n",
        "                    min=2,\n",
        "                    max=6,\n",
        "                    value=4,\n",
        "                    marks={i: str(i) for i in [2, 4, 6]},\n",
        "                ),\n",
        "            ]\n",
        "        ),\n",
        "        dbc.FormGroup(\n",
        "            [\n",
        "                dbc.Spinner(\n",
        "                    [\n",
        "                        dbc.Button(\"Summarize\", id=\"button-run\"),\n",
        "                        html.Div(id=\"time-taken\"),\n",
        "                    ]\n",
        "                )\n",
        "            ]\n",
        "        ),\n",
        "    ],\n",
        "    body=True,\n",
        "    style={\"height\": \"275px\"},\n",
        ")\n",
        "\n",
        "\n",
        "# Define Layout\n",
        "app.layout = dbc.Container(\n",
        "    fluid=True,\n",
        "    children=[\n",
        "        html.H1(\"Dash Automatic Summarization (with DistilBART)\"),\n",
        "        html.Hr(),\n",
        "        dbc.Row(\n",
        "            [\n",
        "                dbc.Col(\n",
        "                    width=5,\n",
        "                    children=[\n",
        "                        controls,\n",
        "                        dbc.Card(\n",
        "                            body=True,\n",
        "                            children=[\n",
        "                                dbc.FormGroup(\n",
        "                                    [\n",
        "                                        dbc.Label(\"Summarized Content\"),\n",
        "                                        dcc.Textarea(\n",
        "                                            id=\"summarized-content\",\n",
        "                                            style={\n",
        "                                                \"width\": \"100%\",\n",
        "                                                \"height\": \"calc(75vh - 275px)\",\n",
        "                                            },\n",
        "                                        ),\n",
        "                                    ]\n",
        "                                )\n",
        "                            ],\n",
        "                        ),\n",
        "                    ],\n",
        "                ),\n",
        "                dbc.Col(\n",
        "                    width=7,\n",
        "                    children=[\n",
        "                        dbc.Card(\n",
        "                            body=True,\n",
        "                            children=[\n",
        "                                dbc.FormGroup(\n",
        "                                    [\n",
        "                                        dbc.Label(\"Original Text (Paste here)\"),\n",
        "                                        dcc.Textarea(\n",
        "                                            id=\"original-text\",\n",
        "                                            style={\"width\": \"100%\", \"height\": \"75vh\"},\n",
        "                                        ),\n",
        "                                    ]\n",
        "                                )\n",
        "                            ],\n",
        "                        )\n",
        "                    ],\n",
        "                ),\n",
        "            ]\n",
        "        ),\n",
        "    ],\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CPsXC0vnpi0a",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "@app.callback(\n",
        "    [Output(\"summarized-content\", \"value\"), Output(\"time-taken\", \"children\")],\n",
        "    [\n",
        "        Input(\"button-run\", \"n_clicks\"),\n",
        "        Input(\"max-length\", \"value\"),\n",
        "        Input(\"num-beams\", \"value\"),\n",
        "    ],\n",
        "    [State(\"original-text\", \"value\")],\n",
        ")\n",
        "def summarize(n_clicks, max_len, num_beams, original_text):\n",
        "    if original_text is None or original_text == \"\":\n",
        "        return \"\", \"Did not run\"\n",
        "\n",
        "    t0 = time.time()\n",
        "\n",
        "    inputs = tokenizer.batch_encode_plus(\n",
        "        [original_text], max_length=1024, return_tensors=\"pt\"\n",
        "    )\n",
        "    inputs = inputs.to(device)\n",
        "\n",
        "    # Generate Summary\n",
        "    summary_ids = model.generate(\n",
        "        inputs[\"input_ids\"],\n",
        "        num_beams=num_beams,\n",
        "        max_length=max_len,\n",
        "        early_stopping=True,\n",
        "    )\n",
        "    out = [\n",
        "        tokenizer.decode(\n",
        "            g, skip_special_tokens=True, clean_up_tokenization_spaces=False\n",
        "        )\n",
        "        for g in summary_ids\n",
        "    ]\n",
        "\n",
        "    t1 = time.time()\n",
        "    time_taken = f\"Summarized on {device} in {t1-t0:.2f}s\"\n",
        "\n",
        "    return out[0], time_taken"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HQUSxYkLsbvp",
        "colab_type": "text"
      },
      "source": [
        "Run the cell below to run your Jupyter Dash app. Click on the **temporary** URL to access the app."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B_6lj6V3pk7I",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "app.run_server(mode='external')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}